# -*- coding: utf-8 -*-
"""
@Time    : 2024/7/12 10:28 
@Author  : ZhangShenao 
@File    : 3.FewShotCoT少样本思维链.py 
@Desc    : 使用Few Shot CoT 少样本思维链模型,进行提示工程
"""
import dotenv
from langchain_community.chat_models import ChatZhipuAI
from langchain_core.output_parsers import StrOutputParser
from langchain_core.prompts import ChatPromptTemplate

dotenv.load_dotenv()

# 设定AI的角色和目标
role_template = '你是一个为花店电商公司工作的AI助手, 你的目标是帮助客户根据他们的喜好做出明智的决定'

# 设定Chain of Thought 提示词
# 这是CoT Prompt的关键部分,需要引导AI进行推理的过程,并且加入一些对话示例(Few-Shot Learning)
cot_prompt = '''
作为一个为花店电商公司工作的AI助手，我的目标是帮助客户根据他们的喜好做出明智的决定。 

我会按部就班的思考，先理解客户的需求，然后考虑各种鲜花的涵义，最后根据这个需求，给出我的推荐。
同时，我也会向客户解释我这样推荐的原因。

示例 1:
  人类：我想找一种象征爱情的花。
  AI：首先，我理解你正在寻找一种可以象征爱情的花。在许多文化中，红玫瑰被视为爱情的象征，这是因为它们的红色通常与热情和浓烈的感情联系在一起。因此，考虑到这一点，我会推荐红玫瑰。红玫瑰不仅能够象征爱情，同时也可以传达出强烈的感情，这是你在寻找的。

示例 2:
  人类：我想要一些独特和奇特的花。
  AI：从你的需求中，我理解你想要的是独一无二和引人注目的花朵。兰花是一种非常独特并且颜色鲜艳的花，它们在世界上的许多地方都被视为奢侈品和美的象征。因此，我建议你考虑兰花。选择兰花可以满足你对独特和奇特的要求，而且，兰花的美丽和它们所代表的力量和奢侈也可能会吸引你。
'''

# 编排Prompt
prompt = ChatPromptTemplate.from_messages([
    ('system', role_template),
    ('system', cot_prompt),
    ('human', '{query}')
])

# 构建LLM
# 设置temperature=0可以让模型生成更确定性的回答,即输出更倾向于最可能的结果
llm = ChatZhipuAI(model='glm-4-air', temperature=0)

# 创建OutputParser
parser = StrOutputParser()

# 构造Chain
chain = prompt | llm | parser

# 执行Chain,获取结果
output = chain.invoke({'query': '我想在母亲节给我母亲送一束花，她比较喜欢蓝色和紫色，你有什么建议吗？'})

# 打印结果
print(output)
